#include<vector>
#include<iostream>
#include<cassert>

#include <formic/lattice_ed/ham_actor_hubb.h>
#include <formic/lattice_ed/indexing.h>
#include <formic/lapack/interface.h>
#include <formic/numeric/numeric.h>
#include <formic/output/output.h>
#include <formic/exception.h>

//-----------------------------------------------------------------------------------
// formic::led::HamActorHubb::HamActorHubb -- constructor
//-----------------------------------------------------------------------------------
template <class SCALAR>
formic::led::HamActorHubb<SCALAR>::HamActorHubb(const int ns, const int na, const int nb,
                                                const std::vector<double> & uvec,
                                                const std::vector<int> & bonds,
                                                const std::vector<double> & bond_strengths,
                                                const int nthread,
                                                const double thresh)
  : _cps(2),
    _nthread(nthread),
    _ns(ns),
    _na(na),
    _nb(nb),
    _thresh(thresh),
    _z_array_a ( formic::led::z_array_size(ns)    , 0 ),
    _z_lookup_a( (ns+1)*(ns+1)                    , 0 ),
    _offsets_a ( (na+1)*(ns+1)*(ns+1)*_cps        , 0 ),
    _spacing_a ( (na+1)*(ns+1)*_cps               , 0 ),
    _z_array_b ( formic::led::z_array_size(ns)    , 0 ),
    _z_lookup_b( (ns+1)*(ns+1)                    , 0 ),
    _offsets_b ( (nb+1)*(ns+1)*(ns+1)*_cps        , 0 ),
    _spacing_b ( (nb+1)*(ns+1)*_cps               , 0 ),
    _work(5*_cps + 2, 0),
    _uvec(uvec),
    _bonds(bonds),
    _bond_strengths(bond_strengths),
    _start_configs_a(nthread, formic::led::Configuration(ns, na, _cps)),
    _start_configs_b(nthread, formic::led::Configuration(ns, nb, _cps)),
    _end_cmpds_a(nthread, -1),
    _end_cmpds_b(nthread, -1),
    _thread_work_ptrs(nthread),
    _thread_workspace()
{

  // prepare alpha and beta indexing arrays
  formic::led::prepare_z_array(ns, &_z_array_a.at(0), &_z_lookup_a.at(0));
  formic::led::prepare_spacing_and_offsets(ns, _cps, na, &_offsets_a.at(0), &_spacing_a.at(0), &_work.at(0));
  formic::led::prepare_z_array(ns, &_z_array_b.at(0), &_z_lookup_b.at(0));
  formic::led::prepare_spacing_and_offsets(ns, _cps, nb, &_offsets_b.at(0), &_spacing_b.at(0), &_work.at(0));

  // get the number of configurations
  const int nconfig_a = formic::n_integer_solutions(ns, na, _cps-1, &_work.at(0));
  const int nconfig_b = formic::n_integer_solutions(ns, nb, _cps-1, &_work.at(0));
  const int nconfig = nconfig_a * nconfig_b;

  // get the number of configurations per thread
  const int ncpt_a = nconfig_a / nthread;
  const int ncpt_b = nconfig_b / nthread;

  // get space for threads to store their vectors
  const size_t buffer_size = 1024;
  _thread_workspace.assign( size_t(nthread) * size_t(nconfig) + size_t(nthread+1) * buffer_size, formic::zero(SCALAR()) );
  for (int i = 0; i < nthread; i++)
    _thread_work_ptrs.at(i) = &_thread_workspace.at( buffer_size * size_t(i+1) + size_t(nconfig) * size_t(i) );

  // get the alpha starting configurations and ending compound indices for each thread
  {
    int count = 0;
    int tnum = 0;
    formic::led::Configuration config(ns, na, _cps);
    _start_configs_a.at(0) = config;
    for ( ; !config.finished(); config++, count++) {
      if (count == ncpt_a && tnum < nthread-1) {
        _end_cmpds_a.at(tnum++) = count;
        _start_configs_a.at(tnum) = config;
        count = 0;
      }
    }
    if (tnum != nthread-1)
      throw formic::Exception("bad tnum, tnum = %i") % tnum;
    _end_cmpds_a.at(tnum) = count;
  }

  // get the beta starting configurations and ending compound indices for each thread
  {
    int count = 0;
    int tnum = 0;
    formic::led::Configuration config(ns, nb, _cps);
    _start_configs_b.at(0) = config;
    for ( ; !config.finished(); config++, count++) {
      if (count == ncpt_b && tnum < nthread-1) {
        _end_cmpds_b.at(tnum++) = count;
        _start_configs_b.at(tnum) = config;
        count = 0;
      }
    }
    if (tnum != nthread-1)
      throw formic::Exception("bad tnum, tnum = %i") % tnum;
    _end_cmpds_b.at(tnum) = count;
  }

}

template formic::led::HamActorHubb<double>::HamActorHubb(const int, const int, const int,
                                                         const std::vector<double> &,
                                                         const std::vector<int> &,
                                                         const std::vector<double> &,
                                                         const int,
                                                         const double);
template formic::led::HamActorHubb<std::complex<double> >::HamActorHubb(const int, const int, const int,
                                                                        const std::vector<double> &,
                                                                        const std::vector<int> &,
                                                                        const std::vector<double> &,
                                                                        const int,
                                                                        const double);

//-----------------------------------------------------------------------------------
// formic::led::HamActorHubb::operate -- function to perform y = A d
//-----------------------------------------------------------------------------------

template <class SCALAR>
void formic::led::HamActorHubb<SCALAR>::operate(const int n, SCALAR * const d, SCALAR * const y, const int myrank, const int tid) {

  formic::xscal(n, formic::zero(SCALAR()), y, 1);

  // create thread attribute object to ensure threads are joinable
  pthread_attr_t attr;
  pthread_attr_init(&attr);
  pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_JOINABLE);

  // create mutexes
  pthread_mutex_t mtx;
  pthread_mutex_init(&mtx, NULL);

  // start threads
  std::vector<pthread_t> threads(_nthread);
  std::vector<typename formic::led::HamActorHubb<SCALAR>::ThreadArgs> thread_args(_nthread);
  for (int i = 0; i < _nthread; i++) {
    thread_args[i].tid = i;
    thread_args[i].obj_ptr = this;
    thread_args[i].mtx_ptr = &mtx;
    thread_args[i].n = n;
    thread_args[i].d = d;
    thread_args[i].y = y;
    thread_args[i].myrank = myrank;
    const int rc = pthread_create(&threads.at(i), &attr, formic::led::thread_func_hubb<SCALAR>, (void *)&thread_args[i]);
    if (rc != 0)
      throw formic::Exception("error in formic::led::HamActorHubb::operate.  pthread_create failed for thread %i with return code %i") % i % rc;
  }

  // wait for threads to finish
  for (int i = 0; i < _nthread; i++) {
    const int rc = pthread_join(threads.at(i), NULL);
    if (rc != 0)
      throw formic::Exception("error in formic::led::HamActorHubb::operate.  pthread_join failed for thread %i with return code %i") % i % rc;
  }

  // destroy mutexes
  pthread_mutex_destroy(&mtx);

  // destroy thread attribute object
  pthread_attr_destroy(&attr);

}

template void formic::led::HamActorHubb<double>::operate(const int n, double * const d, double * const y, const int myrank, const int tid);
template void formic::led::HamActorHubb<std::complex<double> >::operate(const int n,
                                                                        std::complex<double> * const d,
                                                                        std::complex<double> * const y,
                                                                        const int myrank,
                                                                        const int tid);

//-----------------------------------------------------------------------------------
// formic::led::thread_func_hubb -- function to perform one thread's compoenent
//                                  of the Hubbard model Hamiltonian vector product
//-----------------------------------------------------------------------------------

template <class SCALAR> void * formic::led::thread_func_hubb(void * arg) {

  // class used to apply raising/lowering operators
  class RL_ops {

    private:

      // gets the sign associated with a hopping operator
      double get_sign(const formic::led::Configuration & config, const int p, const int q) {
        const int maxpq = std::max(p,q);
        const int minpq = std::min(p,q);
        int retval = 1;
        for (int i = minpq+1; i < maxpq; i++)
          retval = retval - 2 * retval * config[i];
        assert( retval == 1 || retval == -1 );
        return retval;
      }

    public:

      // applies this thread's operators
      void operator()(const int ns,
                      const int cps,
                      const int to,
                      const std::vector<int> & bonds,
                      const std::vector<double> & bond_strengths,
                      const formic::led::Configuration & start_config,
                      const int end_cmpd,
                      const int length,
                      const int multiplier,
                      const int stride,
                      const std::vector<int> & offsets,
                      const std::vector<int> & spacing,
                      const std::vector<int> & z_lookup,
                      const std::vector<int> & z_array,
                      std::vector<int> & work,
                      const SCALAR * const d,
                      SCALAR * const my_y)
      {

        int count = 0;
        for (formic::led::Configuration config = start_config; true; config++, count++) {

          if (count == end_cmpd) break;

          // get compound index
          const int cmpd = formic::led::index_from_occ_vec(ns, cps, to, &offsets[0], &spacing[0],
                                                           &z_lookup[0], &z_array[0], config.c_ptr(), &work[0]);
          const int startd = multiplier * cmpd;

          // apply the bond operators
          for (int i = 0; i < bond_strengths.size(); i++) {

            const int p = bonds.at(2*i+0);
            const int q = bonds.at(2*i+1);
          
            // p^+ q
            config[p]++;
            config[q]--;
            if ( config[p] < cps && config[q] >= 0 ) {
              const SCALAR s = formic::unity(SCALAR()) * this->get_sign(config, p, q);
              const int cmpd2 = formic::led::index_from_occ_vec(ns, cps, to, &offsets[0], &spacing[0],
                                                                &z_lookup[0], &z_array[0], config.c_ptr(), &work[0]);
              const int starty = multiplier * cmpd2;
              formic::xaxpy(length, s * bond_strengths.at(i), d + startd, stride, my_y + starty, stride);
            }
            config[p]--;
            config[q]++;
          
            // q^+ p
            config[q]++;
            config[p]--;
            if ( config[q] < cps && config[p] >= 0 ) {
              const SCALAR s = formic::unity(SCALAR()) * this->get_sign(config, p, q);
              const int cmpd2 = formic::led::index_from_occ_vec(ns, cps, to, &offsets[0], &spacing[0],
                                                                &z_lookup[0], &z_array[0], config.c_ptr(), &work[0]);
              const int starty = multiplier * cmpd2;
              formic::xaxpy(length, s * bond_strengths.at(i), d + startd, stride, my_y + starty, stride);
            }
            config[q]--;
            config[p]++;

          }

        }

      }

  };

  // object to apply hopping operators
  RL_ops rl_ops;

  typedef typename formic::led::HamActorHubb<SCALAR>::ThreadArgs * ArgPtrType;

  ArgPtrType ap = (ArgPtrType)arg;

  const int tid = ap->tid;
  const int myrank = ap->myrank;
  formic::led::HamActorHubb<SCALAR> & obj = *(ap->obj_ptr);
  pthread_mutex_t * const mtx_ptr = ap->mtx_ptr;
  const int n = ap->n;;
  const SCALAR * const d = ap->d;
  SCALAR * const y = ap->y;

  // get and set to zero this thread's scalar work space
  SCALAR * const my_y = obj._thread_work_ptrs.at(tid);
  formic::xscal(n, formic::zero(SCALAR()), my_y, 1);

  // get integer workspace
  std::vector<int> work( 5 * obj._cps + 2, 0);

  // get number of alpha and beta configurations
  const int nconfig_a = formic::n_integer_solutions(obj._ns, obj._na, obj._cps-1);
  const int nconfig_b = formic::n_integer_solutions(obj._ns, obj._nb, obj._cps-1);

  // apply alpha hopping operators
  rl_ops(obj._ns, obj._cps, obj._na, obj._bonds, obj._bond_strengths, obj._start_configs_a.at(tid), obj._end_cmpds_a.at(tid),
         nconfig_b, nconfig_b, 1, obj._offsets_a, obj._spacing_a, obj._z_lookup_a, obj._z_array_a, work, d, my_y);

  // apply beta hopping operators
  rl_ops(obj._ns, obj._cps, obj._nb, obj._bonds, obj._bond_strengths, obj._start_configs_b.at(tid), obj._end_cmpds_b.at(tid),
         nconfig_a, 1, nconfig_b, obj._offsets_b, obj._spacing_b, obj._z_lookup_b, obj._z_array_b, work, d, my_y);

  // apply on site operators
  {

    // create an alpha configuration
    formic::led::Configuration config_a(obj._ns, obj._na, obj._cps);

    // loop over this thread's alpha configurations
    int count = 0;
    for (config_a = obj._start_configs_a.at(tid); true; config_a++, count++) {

      if (count == obj._end_cmpds_a[tid]) break;

      // get alpha compound index
      const int cmpd_a = formic::led::index_from_occ_vec(obj._ns, obj._cps, obj._na, &obj._offsets_a[0], &obj._spacing_a[0],
                                                         &obj._z_lookup_a[0], &obj._z_array_a[0], config_a.c_ptr(), &work[0]);

      // loop over beta configurations
      for (formic::led::Configuration config_b(obj._ns, obj._nb, obj._cps); !config_b.finished(); config_b++) {

        // get beta compound index
        const int cmpd_b = formic::led::index_from_occ_vec(obj._ns, obj._cps, obj._nb, &obj._offsets_b[0], &obj._spacing_b[0],
                                                           &obj._z_lookup_b[0], &obj._z_array_b[0], config_b.c_ptr(), &work[0]);

        // get overall compount index
        const int cmpd = cmpd_a * nconfig_b + cmpd_b;

        // compute the on-site contribution from this configuration
        assert( obj._uvec.size() == obj._ns );
        double u_sum = 0.0;
        for (int i = 0; i < obj._ns; i++)
          u_sum += obj._uvec[i] * config_a[i] * config_b[i];
        my_y[cmpd] += u_sum * d[cmpd];

      }

    }

  }

  // add the results to the total
  pthread_mutex_lock(mtx_ptr);
  formic::xaxpy(n, formic::unity(SCALAR()), my_y, 1, y, 1);
  pthread_mutex_unlock(mtx_ptr);

}

template void * formic::led::thread_func_hubb<double>(void *);
template void * formic::led::thread_func_hubb<std::complex<double> >(void *);
